#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import json
import time
import asyncio
import logging
import argparse
import fractions
from pathlib import Path
from typing import Dict, Set, List, Optional

import cv2
import numpy as np
from aiohttp import web
import aiohttp_cors
import aiortc.sdp
from aiortc import RTCPeerConnection, RTCSessionDescription, VideoStreamTrack, RTCIceCandidate

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[
        logging.StreamHandler()
    ]
)
logger = logging.getLogger("webrtc_server")

# 当前目录
BASE_DIR = Path(os.path.dirname(os.path.abspath(__file__)))

# 图片目录
IMAGE_DIR = BASE_DIR / "static" / "img"

# 房间管理
class Room:
    def __init__(self, room_id: str):
        self.room_id = room_id
        self.clients: Set[web.WebSocketResponse] = set()
        self.server_pc: Optional[RTCPeerConnection] = None
        
    def add_client(self, ws: web.WebSocketResponse):
        self.clients.add(ws)
        
    def remove_client(self, ws: web.WebSocketResponse):
        if ws in self.clients:
            self.clients.remove(ws)

# 房间字典
rooms: Dict[str, Room] = {}

# 错误处理函数
async def handle_error(ws, error_message, exception=None):
    """处理错误并发送错误消息给客户端"""
    error_info = str(exception) if exception else "未知错误"
    logger.error(f"{error_message}: {error_info}")
    
    try:
        await ws.send_json({
            "type": "error",
            "message": f"{error_message}: {error_info}"
        })
    except Exception as e:
        logger.error(f"发送错误消息失败: {e}")


# 自定义视频轨道 - 屏幕捕获
class ScreenCaptureVideoStreamTrack(VideoStreamTrack):
    def __init__(self):
        super().__init__()
        self.counter = 0
        self.frame_interval = 1.0 / 30.0  # 30 FPS
        self.last_frame_time = time.time()
        self.images = []  # 图像列表
        self.load_images()
    
    def load_images(self):
        """加载图像文件"""
        try:
            # 获取图片目录中的所有jpg文件
            jpg_files = sorted(IMAGE_DIR.glob("*.jpg"))
            
            if not jpg_files:
                logger.warning(f"在目录 {IMAGE_DIR} 中未找到jpg图片文件")
                # 创建一个默认的黑色图像
                default_frame = np.zeros((480, 640, 3), dtype=np.uint8)
                self.images = [default_frame]
                return
            
            logger.info(f"找到 {len(jpg_files)} 个图片文件")
            
            # 加载所有图片
            for jpg_file in jpg_files:
                try:
                    # 使用cv2加载图片
                    img = cv2.imread(str(jpg_file))
                    if img is not None:
                        # 转换BGR到RGB (cv2默认是BGR，但WebRTC需要RGB)
                        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                        self.images.append(img_rgb)
                        logger.debug(f"成功加载图片: {jpg_file.name}, 尺寸: {img_rgb.shape}")
                    else:
                        logger.warning(f"无法加载图片: {jpg_file}")
                except Exception as e:
                    logger.error(f"加载图片 {jpg_file} 时出错: {e}")
            
            if not self.images:
                logger.warning("没有成功加载任何图片，使用默认黑色图像")
                default_frame = np.zeros((480, 640, 3), dtype=np.uint8)
                self.images = [default_frame]
            else:
                logger.info(f"成功加载 {len(self.images)} 张图片")
                
        except Exception as e:
            logger.error(f"加载图片时发生错误: {e}")
            # 创建默认图像作为备用
            default_frame = np.zeros((480, 640, 3), dtype=np.uint8)
            self.images = [default_frame]
        
    async def recv(self):
        """生成视频帧"""
        try:
            # 控制帧率
            current_time = time.time()
            time_since_last_frame = current_time - self.last_frame_time
            
            if time_since_last_frame < self.frame_interval:
                # 等待到下一帧时间
                await asyncio.sleep(self.frame_interval - time_since_last_frame)
            
            self.last_frame_time = time.time()
            
            # 如果没有图片，返回空帧
            if not self.images:
                logger.warning("没有可用的图片，返回黑色帧")
                frame = np.zeros((480, 640, 3), dtype=np.uint8)
            else:
                # 循环使用图片
                image_index = self.counter % len(self.images)
                frame = self.images[image_index].copy()
            
            # 增加计数器
            self.counter += 1
            
            # 创建VideoFrame
            from av import VideoFrame
            av_frame = VideoFrame.from_ndarray(frame, format="rgb24")
            
            # 设置时间戳
            av_frame.pts = self.counter
            av_frame.time_base = fractions.Fraction(1, 30)  # 30 FPS
            
            return av_frame
            
        except Exception as e:
            logger.error(f"生成视频帧时出错: {e}")
            # 返回黑色帧作为备用
            frame = np.zeros((480, 640, 3), dtype=np.uint8)
            from av import VideoFrame
            av_frame = VideoFrame.from_ndarray(frame, format="rgb24")
            av_frame.pts = self.counter
            av_frame.time_base = fractions.Fraction(1, 30)
            self.counter += 1
            return av_frame

# 信令处理
async def websocket_handler(request):
    ws = web.WebSocketResponse()
    await ws.prepare(request)
    
    client_id = str(id(ws))
    current_room = None
    
    logger.info(f"客户端 {client_id} 连接")
    
    try:
        async for mesg in ws:
            if mesg.type == web.WSMsgType.TEXT:
                data = json.loads(mesg.data)
                logger.debug(f"收到消息: {data}")
                
                if data.get("type") == "join":
                    # 加入房间
                    room_id = data.get("room", "default")
                    
                    # 如果房间不存在，创建它
                    if room_id not in rooms:
                        rooms[room_id] = Room(room_id)
                    
                    # 将客户端添加到房间
                    current_room = rooms[room_id]
                    current_room.add_client(ws)
                    
                    # 通知客户端已准备好
                    await ws.send_json({"type": "ready"})
                    logger.info(f"客户端 {client_id} 加入房间 {room_id}")
                    
                elif data.get("type") == "offer" and current_room:
                    # 处理WebRTC offer
                    try:
                        offer = data.get("sdp")
                        if not offer:
                            await handle_error(ws, "无效的offer格式", None)
                            return
                            
                        # 创建对等连接
                        if not current_room.server_pc:
                            # 创建RTCPeerConnection
                            from aiortc import RTCConfiguration, RTCIceServer
                            configuration = RTCConfiguration(
                                iceServers=[RTCIceServer(urls=["stun:stun.l.google.com:19302"])]
                            )
                            current_room.server_pc = RTCPeerConnection(configuration)
                            
                            # 先添加视频轨道（必须在设置远程描述之前）
                            video_track = ScreenCaptureVideoStreamTrack()
                            logger.info(f"创建视频轨道: {video_track}")
                            current_room.server_pc.addTrack(video_track)
                            logger.info("视频轨道已添加到PeerConnection")
                            
                            # 然后设置远程描述
                            sdp_text = offer.get("sdp")
                            sdp_type = offer.get("type")
                            
                            if not sdp_text or not sdp_type:
                                await handle_error(ws, "SDP offer缺少必要字段", None)
                                return
                                
                            # 调试：打印收到的offer SDP
                            logger.debug(f"收到的offer SDP: {sdp_text}")
                            
                            await current_room.server_pc.setRemoteDescription(RTCSessionDescription(
                                sdp=sdp_text,
                                type=sdp_type
                            ))
                            
                            @current_room.server_pc.on("icecandidate")
                            async def on_icecandidate(candidate):
                                if candidate and ws.closed is False:
                                    logger.info(f"发送ICE候选者: {candidate}")
                                    await ws.send_json({
                                        "type": "candidate",
                                        "candidate": candidate.asDict()
                                    })
                                    logger.info("ICE候选者已发送给客户端")
                                else:
                                    logger.debug(f"跳过ICE候选者发送: candidate={candidate}, ws_closed={ws.closed}")
                            
                            @current_room.server_pc.on("connectionstatechange")
                            async def on_connectionstatechange():
                                if current_room.server_pc:
                                    logger.info(f"连接状态变化: {current_room.server_pc.connectionState}")
                                else:
                                    logger.warning("连接状态变化回调中server_pc为None")
                            
                            @current_room.server_pc.on("track")
                            def on_track(track):
                                logger.info(f"收到轨道: {track.kind}")
                        
                        # 远程描述已在上面设置
                        
                        # 创建应答
                        answer = await current_room.server_pc.createAnswer()
                        
                        # 调试：打印生成的SDP
                        logger.debug(f"生成的answer SDP: {answer.sdp}")
                        
                        await current_room.server_pc.setLocalDescription(answer)
                        
                        # 获取本地描述
                        answer_dict = {
                            "type": answer.type,
                            "sdp": answer.sdp
                        }
                        
                        # 调试：打印发送的answer
                        logger.debug(f"发送的answer内容: {answer_dict}")
                            
                        # 发送应答
                        logger.info(f"发送answer给客户端: {answer_dict['type']}")
                        await ws.send_json({
                            "type": "answer",
                            "sdp": answer_dict
                        })
                        logger.info("answer已发送给客户端")
                    except (KeyError, TypeError) as e:
                        await handle_error(ws, "解析SDP offer失败", e)
                        if current_room and current_room.server_pc:
                            await current_room.server_pc.close()
                            current_room.server_pc = None
                    except Exception as e:
                        await handle_error(ws, "处理offer失败", e)
                        if current_room and current_room.server_pc:
                            await current_room.server_pc.close()
                            current_room.server_pc = None
                    
                elif data.get("type") == "candidate" and current_room and current_room.server_pc:
                    candidate_data = data.get("candidate")
                    if candidate_data and candidate_data.get("candidate"):
                        try:
                            # 从SDP字符串创建候选者对象
                            candidate = aiortc.sdp.candidate_from_sdp(candidate_data["candidate"])
                            # 设置sdpMid和sdpMLineIndex
                            candidate.sdpMid = candidate_data["sdpMid"]
                            candidate.sdpMLineIndex = candidate_data["sdpMLineIndex"]
                            
                            # 添加ICE候选者
                            await current_room.server_pc.addIceCandidate(candidate)
                            logger.info(f"成功添加客户端ICE候选者: {candidate.candidate}")
                        except Exception as e:
                            await handle_error(ws, "处理ICE候选者失败", e)
                    else:
                        logger.warning("收到空的或无效的ICE候选者")
                        
                elif data.get("type") == "leave" and current_room:
                    # 离开房间
                    if current_room.server_pc:
                        await current_room.server_pc.close()
                        current_room.server_pc = None
                    
                    current_room.remove_client(ws)
                    logger.info(f"客户端 {client_id} 离开房间 {current_room.room_id}")
                    current_room = None
            
            elif mesg.type == web.WSMsgType.ERROR:
                logger.error(f"WebSocket连接错误: {ws.exception()}")
    
    finally:
        # 清理
        if current_room:
            if current_room.server_pc:
                await current_room.server_pc.close()
                current_room.server_pc = None
            
            current_room.remove_client(ws)
            logger.info(f"客户端 {client_id} 断开连接，离开房间 {current_room.room_id}")
            
            # 如果房间为空，删除它
            if not current_room.clients:
                del rooms[current_room.room_id]
                logger.info(f"房间 {current_room.room_id} 已删除")
        
        logger.info(f"客户端 {client_id} 断开连接")
    
    return ws

# 主页路由
async def index_handler(request):
    with open(BASE_DIR / "templates" / "webrtc.html") as f:
        return web.Response(text=f.read(), content_type="text/html")

# 创建应用
async def create_app():
    app = web.Application()
    
    # 设置路由
    app.router.add_get("/", index_handler)
    app.router.add_get("/signaling", websocket_handler)
    
    # 静态文件
    app.router.add_static("/static", BASE_DIR / "static")
    
    # 设置CORS
    cors = aiohttp_cors.setup(app, defaults={
        "*": aiohttp_cors.ResourceOptions(
            allow_credentials=True,
            expose_headers="*",
            allow_headers="*"
        )
    })
    
    # 为所有路由应用CORS
    for route in list(app.router.routes()):
        if not isinstance(route.resource, web.StaticResource):
            cors.add(route)
    
    return app

# 主函数
async def main():
    parser = argparse.ArgumentParser(description="WebRTC远程桌面服务器")
    parser.add_argument("--host", default="0.0.0.0", help="监听主机")
    parser.add_argument("--port", type=int, default=8080, help="监听端口")
    parser.add_argument("--verbose", action="store_true", help="启用详细日志")
    
    args = parser.parse_args()
    
    if args.verbose:
        logger.setLevel(logging.DEBUG)
    
    app = await create_app()
    
    # 启动服务器
    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, args.host, args.port)
    
    logger.info(f"启动WebRTC远程桌面服务器 http://{args.host}:{args.port}")
    await site.start()
    
    # 保持运行
    while True:
        await asyncio.sleep(3600)  # 每小时检查一次

# 入口点
if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("服务器已停止")